import time
import logging

from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware


# 设置日志输出
logger = logging.getLogger("request_logger")
logger.setLevel(logging.INFO)

handler = logging.StreamHandler()
formatter = logging.Formatter('[%(asctime)s] %(levelname)s: %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)


class RequestTimerMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        start_time = time.time()

        response = await call_next(request)     # 执行请求处理链

        process_time = time.time() - start_time

        formatted_time = f"{process_time * 100:.2f}ms"

        logger.info(f"{request.method} {request.url.path} - 耗时: {formatted_time}")

        # 可以添加到响应头中返回
        response.headers["X-Process-Time"] = formatted_time
        return response